#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
按「分批次读取 → 筛选 → 保存」流程处理超大文本：
- 输入: data/train_processed.txt (CSV，含列: fact, accusation)
- 映射: data/class.json (罪名 -> 大类)
- 输出: data/train_processed.json (单个JSON数组，逐条流式写入，避免高内存)

筛选规则:
- 解析每行的 accusation 列(形如 ['盗窃', '诈骗'] 的字符串)，
- 使用 class.json 做重新划分(将罪名映射为对应的大类)，
- 若该行至少能映射出一个大类则保留，否则丢弃。

内存控制:
- 流式逐行读取，不整体加载。
- 按 batch_size 条写入一次(实际逐条写入，这里主要用于定期flush与进度输出)。

可通过 --max-rows 做小样本验证。
"""

import argparse
import ast
import csv
import json
import os
import re
import sys
from typing import Dict, List


def load_class_mapping(path: str) -> Dict[str, str]:
    """读取罪名->大类映射。"""
    with open(path, 'r', encoding='utf-8') as f:
        text = f.read()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        # 处理可能的 BOM 或异常空白
        return json.loads(text.strip('\ufeff').strip())


def parse_accusation_field(raw: str) -> List[str]:
    """将 CSV 第三列 accusation 的字符串解析为罪名列表。
    兼容以下格式:
    - Python风格: "['盗窃', '诈骗']"
    - JSON风格:  "[\"盗窃\", \"诈骗\"]"
    - 单个字符串: "盗窃"
    解析失败时，尝试按中英文逗号切分。
    """
    if raw is None:
        return []
    s = raw.strip()
    if not s:
        return []
    # 首选安全的 literal_eval
    try:
        val = ast.literal_eval(s)
        if isinstance(val, list):
            return [str(x).strip() for x in val if str(x).strip()]
        return [str(val).strip()] if str(val).strip() else []
    except Exception:
        pass
    # 尝试 JSON
    try:
        val = json.loads(s)
        if isinstance(val, list):
            return [str(x).strip() for x in val if str(x).strip()]
        return [str(val).strip()] if str(val).strip() else []
    except Exception:
        pass
    # 兜底: 按中英文逗号拆分，并去除引号
    parts = [p.strip().strip("'").strip('"') for p in re.split(r'[，,]+', s)]
    return [p for p in parts if p]


def process(
    input_path: str,
    class_json_path: str,
    output_path: str,
    batch_size: int = 50000,
    max_rows: int | None = None,
) -> None:
    mapping = load_class_mapping(class_json_path)

    total = 0
    kept = 0
    dropped = 0

    # 确保输出目录存在
    os.makedirs(os.path.dirname(output_path), exist_ok=True)

    # 注意: 使用 utf-8-sig 以自动处理可能的 BOM
    with open(input_path, 'r', encoding='utf-8-sig', newline='') as fin, \
         open(output_path, 'w', encoding='utf-8') as fout:
        reader = csv.DictReader(fin)

        # 写 JSON 数组的起始标记
        fout.write('[\n')
        first_written = False

        for row in reader:
            total += 1
            if max_rows is not None and total > max_rows:
                break

            fact = row.get('fact')
            acc_raw = row.get('accusation')

            # 清洗 fact（若为空或缺失则丢弃）
            if not fact or not isinstance(fact, str):
                dropped += 1
                continue

            accusations = parse_accusation_field(acc_raw)
            if not accusations:
                dropped += 1
                continue

            categories = sorted({mapping[a] for a in accusations if a in mapping})
            # 至少映射出一个大类才保留
            if not categories:
                dropped += 1
                continue

            obj = {
                'fact': fact,
                'accusation': accusations,
                'categories': categories,
            }

            line = json.dumps(obj, ensure_ascii=False)
            if first_written:
                fout.write(',\n')
            fout.write(line)
            first_written = True
            kept += 1

            # 批次 flush + 进度
            if batch_size and (kept % batch_size == 0):
                fout.flush()
                print(
                    f"进度: 处理 {total:,} 行, 保留 {kept:,} 行, 丢弃 {dropped:,} 行...",
                    file=sys.stderr,
                )

        # 写 JSON 数组的结束标记
        fout.write('\n]\n')

    # 最终统计输出到 stderr，避免干扰主输出文件
    print(
        f"完成: 总行数 {total:,}, 保留 {kept:,}, 丢弃 {dropped:,}. 输出 => {output_path}",
        file=sys.stderr,
    )


def main() -> None:
    default_input = r'D:\LearnPython\crime-type-discrimination\data\train_processed.txt'
    default_class = r'D:\LearnPython\crime-type-discrimination\data\class.json'
    default_output = r'D:\LearnPython\crime-type-discrimination\data\train_processed.json'

    parser = argparse.ArgumentParser(description='按批处理并重划分类别，控制内存占用')
    parser.add_argument('--input', type=str, default=default_input, help='输入CSV路径')
    parser.add_argument('--class-json', type=str, default=default_class, help='罪名->大类 映射JSON路径')
    parser.add_argument('--output', type=str, default=default_output, help='输出JSON路径')
    parser.add_argument('--batch-size', type=int, default=50000, help='每批处理/flush条数')
    parser.add_argument('--max-rows', type=int, default=None, help='仅处理前N行(用于小样本验证)')

    args = parser.parse_args()

    process(
        input_path=args.input,
        class_json_path=args.class_json,
        output_path=args.output,
        batch_size=args.batch_size,
        max_rows=args.max_rows,
    )


if __name__ == '__main__':
    main()